import gensim

class Vectorization:

    def __init__(self):
        pass

    @staticmethod
    def word2Vec_train(good_doc, doc_tag):
        pre_doc = []
    # 使用count当做每个句子的“标签”，标签和每个句子是一一对应的
        count_tag = [count for count in range(len(good_doc))]
        TD = gensim.models.doc2vec.TaggedDocument
        for (i, doc) in zip(count_tag, good_doc):
          doc_new = TD(doc, tags=[i])
          pre_doc.append(doc_new)

        model = gensim.models.Doc2Vec(pre_doc, vector_size=50, window=3,
                                  min_count=1, workers=4)
        return model


    """
    使用doc2vec 对训练文档进行训练 
    """
    @staticmethod
    def word2Vec_train(good_doc):
        pre_doc = []
        # 使用count当做每个句子的“标签”，标签和每个句子是一一对应的
        count_tag = [count for count in range(len(good_doc))]
        TD = gensim.models.doc2vec.TaggedDocument
        for (i, doc) in zip(count_tag, good_doc):
            doc_new = TD(doc, tags=[i])
            pre_doc.append(doc_new)
            #pre_doc, vector_size=20, window=5, min_count=1, workers=6, epochs=25
         #pre_doc, vector_size=15, window=5, min_count=2, workers=6, epochs=30
        model = gensim.models.Doc2Vec(pre_doc, vector_size=20, window=5,
                                      min_count=1, workers=6, epochs=35)
        return model


    """
    将训练好的文本向量化模型应用于测试数据 返回向量化结果
    """
    @staticmethod
    def vectorizationWithModel(model_path, data_list):
        Model = gensim.models.doc2vec.Doc2Vec.load(model_path)
        result = []
        for sentence in data_list:
            vector = Model.infer_vector(sentence)
            result.append(vector)
        return result

    """
    向量化单个句子
    """
    @staticmethod
    def vectorizationWithModelOnSentence(cutSentence,model_path):
        Model = gensim.models.doc2vec.Doc2Vec.load(model_path)
        vector = Model.infer_vector(cutSentence)
        return vector



